#!/bin/bash

# 渐进式压测脚本
# 避免瞬时并发冲击导致 WSL 断开连接

echo "======================================"
echo "渐进式压测工具"
echo "======================================"

# 检查客户端程序是否存在
if [ ! -f "build/bin/client_test" ]; then
    echo "错误: 找不到 build/bin/client_test"
    echo "请先编译项目"
    exit 1
fi

# 检查服务器是否运行
if ! pgrep -x "server_test" > /dev/null; then
    echo "错误: 服务器未运行"
    echo "请先在另一个终端启动服务器: ./build/bin/server_test"
    exit 1
fi

SERVER_PID=$(pgrep -x "server_test")
echo "找到服务器进程 PID: $SERVER_PID"

# 配置参数
read -p "启动多少个客户端? [500]: " CLIENT_COUNT
CLIENT_COUNT=${CLIENT_COUNT:-500}

read -p "每次启动间隔（毫秒）? [100]: " INTERVAL_MS
INTERVAL_MS=${INTERVAL_MS:-100}
INTERVAL=$(echo "scale=3; $INTERVAL_MS / 1000" | bc)

read -p "每启动多少个暂停一次? [50]: " PAUSE_EVERY
PAUSE_EVERY=${PAUSE_EVERY:-50}

read -p "暂停时间（秒）? [2]: " PAUSE_DURATION
PAUSE_DURATION=${PAUSE_DURATION:-2}

echo ""
echo "压测配置:"
echo "  客户端总数: $CLIENT_COUNT"
echo "  启动间隔: ${INTERVAL_MS}ms"
echo "  每 $PAUSE_EVERY 个暂停 ${PAUSE_DURATION}s"
echo ""

read -p "按回车开始压测..."

# 创建日志目录
LOG_DIR="stress_test_logs"
mkdir -p $LOG_DIR
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
LOG_FILE="${LOG_DIR}/stress_${TIMESTAMP}.log"

# 记录初始状态
echo "======================================" | tee -a $LOG_FILE
echo "压测开始时间: $(date)" | tee -a $LOG_FILE
echo "初始状态:" | tee -a $LOG_FILE
echo "  FD 数量: $(ls /proc/$SERVER_PID/fd 2>/dev/null | wc -l)" | tee -a $LOG_FILE
echo "  内存使用: $(ps -o rss= -p $SERVER_PID | awk '{print $1/1024 "MB"}')" | tee -a $LOG_FILE
echo "======================================" | tee -a $LOG_FILE

# 启动客户端
CLIENT_PIDS=()
FAILED_COUNT=0

for i in $(seq 1 $CLIENT_COUNT); do
    # 检查服务器是否还活着
    if ! kill -0 $SERVER_PID 2>/dev/null; then
        echo "" | tee -a $LOG_FILE
        echo "⚠️  服务器进程已终止！" | tee -a $LOG_FILE
        echo "已启动客户端数: $((i-1))" | tee -a $LOG_FILE
        break
    fi

    # 启动客户端
    ./build/bin/client_test > /dev/null 2>&1 &
    CLIENT_PID=$!
    CLIENT_PIDS+=($CLIENT_PID)

    # 检查客户端是否成功启动
    sleep 0.01
    if ! kill -0 $CLIENT_PID 2>/dev/null; then
        FAILED_COUNT=$((FAILED_COUNT + 1))
    fi

    # 显示进度
    echo -ne "\r已启动: $i/$CLIENT_COUNT 客户端 (失败: $FAILED_COUNT)"

    # 定期暂停和报告
    if [ $((i % PAUSE_EVERY)) -eq 0 ]; then
        FD_COUNT=$(ls /proc/$SERVER_PID/fd 2>/dev/null | wc -l)
        MEM_MB=$(ps -o rss= -p $SERVER_PID 2>/dev/null | awk '{print $1/1024}')
        TCP_COUNT=$(lsof -p $SERVER_PID 2>/dev/null | grep TCP | grep -v LISTEN | wc -l)

        echo "" | tee -a $LOG_FILE
        echo "[$i/$CLIENT_COUNT] FD: $FD_COUNT | TCP连接: $TCP_COUNT | 内存: ${MEM_MB}MB | 失败: $FAILED_COUNT" | tee -a $LOG_FILE

        sleep $PAUSE_DURATION
    else
        sleep $INTERVAL
    fi
done

echo "" | tee -a $LOG_FILE
echo "" | tee -a $LOG_FILE
echo "======================================" | tee -a $LOG_FILE
echo "所有客户端已启动" | tee -a $LOG_FILE
echo "总数: ${#CLIENT_PIDS[@]} | 失败: $FAILED_COUNT" | tee -a $LOG_FILE
echo "======================================" | tee -a $LOG_FILE

# 监控一段时间
echo "" | tee -a $LOG_FILE
echo "持续监控 30 秒..." | tee -a $LOG_FILE

for i in {1..30}; do
    if ! kill -0 $SERVER_PID 2>/dev/null; then
        echo "⚠️  服务器进程已终止！" | tee -a $LOG_FILE
        break
    fi

    FD_COUNT=$(ls /proc/$SERVER_PID/fd 2>/dev/null | wc -l)
    MEM_MB=$(ps -o rss= -p $SERVER_PID 2>/dev/null | awk '{print $1/1024}')
    TCP_COUNT=$(lsof -p $SERVER_PID 2>/dev/null | grep TCP | grep -v LISTEN | wc -l)
    ACTIVE_CLIENTS=$(ps -ef | grep client_test | grep -v grep | wc -l)

    echo "[${i}s] FD: $FD_COUNT | TCP: $TCP_COUNT | 内存: ${MEM_MB}MB | 活跃客户端: $ACTIVE_CLIENTS" | tee -a $LOG_FILE

    sleep 1
done

# 最终状态
echo "" | tee -a $LOG_FILE
echo "======================================" | tee -a $LOG_FILE
echo "压测结束时间: $(date)" | tee -a $LOG_FILE
echo "最终状态:" | tee -a $LOG_FILE

if kill -0 $SERVER_PID 2>/dev/null; then
    FD_COUNT=$(ls /proc/$SERVER_PID/fd 2>/dev/null | wc -l)
    MEM_MB=$(ps -o rss= -p $SERVER_PID 2>/dev/null | awk '{print $1/1024}')
    TCP_COUNT=$(lsof -p $SERVER_PID 2>/dev/null | grep TCP | grep -v LISTEN | wc -l)
    TIME_WAIT=$(netstat -an 2>/dev/null | grep TIME_WAIT | wc -l)

    echo "  服务器状态: 运行中 ✓" | tee -a $LOG_FILE
    echo "  FD 数量: $FD_COUNT" | tee -a $LOG_FILE
    echo "  TCP 连接: $TCP_COUNT" | tee -a $LOG_FILE
    echo "  内存使用: ${MEM_MB}MB" | tee -a $LOG_FILE
    echo "  TIME_WAIT 连接: $TIME_WAIT" | tee -a $LOG_FILE
else
    echo "  服务器状态: 已终止 ✗" | tee -a $LOG_FILE
fi

echo "======================================" | tee -a $LOG_FILE
echo "" | tee -a $LOG_FILE
echo "完整日志: $LOG_FILE"

# 询问是否清理客户端进程
echo ""
read -p "是否终止所有客户端进程? (y/n): " cleanup

if [ "$cleanup" = "y" ]; then
    echo "正在终止客户端..."
    pkill -9 client_test
    echo "✓ 客户端已清理"
fi

echo ""
echo "压测完成"
